from collections.abc import Iterable, Iterator
from typing import Any, TypeVar, TypedDict

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, SkipExhausted, pause, resume, skip

AnySample = TypeVar("AnySample")


class BoundedIteratorState(TypedDict):
    iterator_state: Any
    remaining_samples_n: int


class BoundedIterator(DatasetIterator[AnySample]):
    def __init__(self, iterator: Iterable[AnySample] | Iterator[AnySample], remaining_samples_n: int):
        self._iterator = iter(iterator)
        self._remaining_samples_n = remaining_samples_n

    @property
    def remaining_samples_n(self) -> int:
        return self._remaining_samples_n

    def __next__(self) -> AnySample:
        if self._remaining_samples_n < 1:
            raise StopIteration
        sample = next(self._iterator)
        self._remaining_samples_n -= 1
        return sample

    def __skip__(self, samples_n: int):
        if self._remaining_samples_n < samples_n:
            skip(self, self._remaining_samples_n)
            raise SkipExhausted(self._remaining_samples_n)

        try:
            skip(self._iterator, samples_n)
            self._remaining_samples_n -= samples_n
        except SkipExhausted as err:
            self._remaining_samples_n -= err.skipped_samples_n

    def __pause__(self) -> BoundedIteratorState:
        return BoundedIteratorState(
            iterator_state=pause(self._iterator),
            remaining_samples_n=self._remaining_samples_n)


class BoundedDataset(Dataset[AnySample]):
    def __init__(self, dataset: Dataset[AnySample], bound_samples_n: int):
        self._dataset = dataset
        self._bound_samples_n = bound_samples_n

    @property
    def dataset(self) -> Dataset[AnySample]:
        return self._dataset

    @property
    def bound_samples_n(self) -> int:
        return self._bound_samples_n

    def __len__(self) -> int:
        return min(len(self._dataset), self._bound_samples_n)

    def __getitem__(self, index: int, /) -> AnySample:
        if index >= self._bound_samples_n:
            raise IndexError(f"Index {index} is out of bounds for dataset of size {self._bound_samples_n}")
        return self._dataset[index]

    def __iter__(self) -> BoundedIterator[AnySample]:
        return BoundedIterator(iter(self._dataset), self._bound_samples_n)

    def __resume__(self, state: BoundedIteratorState | None = None, /) -> BoundedIterator[AnySample]:
        return BoundedIterator(resume(self._dataset, state["iterator_state"]), state["remaining_samples_n"])
